#!/usr/bin/env python3
'''
    主脚本，用于监听无线网络，并进行智能配网
'''
from scapy.all import Dot11, sniff, conf, raw
import os
import sys
from pathlib import Path
import time
import subprocess
import threading
import socket
from datetime import datetime
import json
import logging
from logging.handlers import RotatingFileHandler
import pywifi
from pywifi import const
# 常量定义
INTERFACE = "wlan0mon"
BROADCAST_MAC = "ff:ff:ff:ff:ff:ff"
HEADER_LEN = 361
ASCII_START = 127
ASCII_END = 128


# 全局状态变量
is_active = False
frame_data = []
data_dict = {}
total_len = 0
total_len_list = []
total_len_count = 0
current_index = -1
target_data = []
sniffing_thread = None
sniffing_event = threading.Event()
exit_event = threading.Event()  # 新增退出信号
has_check_repeat_flag = False


def setup_logger():
    logger = logging.getLogger('smart_config')
    
    # 关键修复：检查是否已有handler，避免重复添加
    if logger.handlers:  # 如果已有handler，直接返回现有logger
        return logger

        # 清空现有日志文件
    log_file = '/var/log/smart_config.log'
    if os.path.exists(log_file):
        with open(log_file, 'w', encoding='utf-8') as f:
            f.truncate()

    logger.setLevel(logging.DEBUG)
    formatter = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )

    # 文件处理器（仅添加一次）
    file_handler = RotatingFileHandler(
        log_file,
        maxBytes=10 * 1024 * 1024,
        backupCount=5,
        encoding='utf-8'
    )
    file_handler.setLevel(logging.INFO)
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)

    # 控制台处理器（仅添加一次）
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.DEBUG)
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)

    return logger



def is_wifi_connected():
    '''
    检查是否存在无线网络连接。
    如果没有, 则返回False, 执行无线配网流程
    如果有, 则返回True, 退出程序
    '''
    wifi = pywifi.PyWiFi()
    iface = wifi.interfaces()[0]  # 获取第一个无线网卡
    return iface.status() == const.IFACE_CONNECTED

def check_root():
    if os.geteuid() != 0:
        logger.info("❌ 错误：需要root权限！")
        sys.exit(1)

def ascii_char(code: int) -> str:
    if 32 <= code <= 126:
        return chr(code)
    return f"(非打印:{code})"

def reset():
    global is_active, frame_data, has_check_repeat_flag
    is_active = False
    frame_data.clear()
    has_check_repeat_flag = False

def parse_wifi_credentials(credentials_str):
    try:
        ssid, password = credentials_str.split(':', maxsplit=1)
        return ssid.strip(), password.strip()
    except ValueError:
        raise ValueError("无效的凭证格式，应为 'SSID:password'")

def switch_to_managed(monitor_iface="wlan0mon", managed_iface="wlan0"):
    global INTERFACE
    logger.info(f"📡 停止监控模式接口: {monitor_iface}")
    subprocess.run(
        ["airmon-ng", "stop", monitor_iface],
        check=True,
        capture_output=True,
        text=True
    )
    logger.info(f"🔄 设置接口 {managed_iface} 为 managed 模式")
    subprocess.run(
        ["ip", "link", "set", managed_iface, "down"],
        check=True,
        capture_output=True,
        text=True
    )
    subprocess.run(
        ["iw", managed_iface, "set", "type", "managed"],
        check=True,
        capture_output=True,
        text=True
    )
    subprocess.run(
        ["ip", "link", "set", managed_iface, "up"],
        check=True,
        capture_output=True,
        text=True
    )
    logger.info("🔄 启动 NetworkManager")
    subprocess.run(
        ["systemctl", "start", "NetworkManager"],
        check=True,
        capture_output=True,
        text=True
    )
    time.sleep(5)
    INTERFACE = managed_iface
    logger.info(f"✅ 已切换到 managed 模式: {managed_iface}")
    

def switch_to_monitor(managed_iface="wlan0", monitor_iface="wlan0mon", timeout=10):
    """切换为监控模式（带启动状态检测）"""
    global INTERFACE
    
    logger.info("🔄 停止 NetworkManager")
    subprocess.run(
        ["airmon-ng", "check", "kill"],
        check=True,
        capture_output=True,
        text=True
    )
    time.sleep(3)
    logger.info(f"📡 尝试启动监控模式: {monitor_iface}")
    try:
        # 2. 启动监控模式（指定频道6）
        subprocess.run(
            ["airmon-ng", "start", managed_iface, "6"],
            check=True,
            capture_output=True,
            text=True
        )
        
        # 3. 循环检测接口状态（带超时）
        start_time = time.time()
        while time.time() - start_time < timeout:
            # 检测接口是否存在
            if not os.path.exists(f"/sys/class/net/{monitor_iface}"):
                time.sleep(1)
                continue
            logger.info(f"{monitor_iface}存在！")
            # 检测监控模式是否激活
            result = subprocess.run(["iwconfig", monitor_iface], 
                                  capture_output=True, text=True)
            if "Mode:Monitor" in result.stdout:
                INTERFACE = monitor_iface
                logger.info(f"✅ 监控模式启动成功 | 接口: {monitor_iface} | 频道: 6")
                return True
            time.sleep(1)

        logger.info(f"❌ 超时：{monitor_iface} 未在 {timeout} 秒内进入监控模式")
        return False
        
    except subprocess.CalledProcessError as e:
        logger.info(f"❌ 启动失败: {e.stderr.strip()}")
        return False


def broadcast_ack(port=12345, ack_message="ACK"):
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        sock.settimeout(0.5)
        message = ack_message.encode('utf-8')
        for i in range(10):
            sock.sendto(message, ('255.255.255.255', port))
            logger.info(f"✅ 已广播ACK到255.255.255.255:{port}")
    except Exception as e:
        logger.info(f"❌ 广播失败: {str(e)}")
    finally:
        sock.close()

def start_sniffing():
    global INTERFACE
    while not exit_event.is_set():  # 检查退出信号
        if sniffing_event.is_set():
            try:
                logger.info(f"🚀 开始嗅探接口: {INTERFACE}")
                sniff(iface=INTERFACE, prn=handle_packet, store=0, monitor=True, timeout=1)
            except Exception as e:
                logger.info(f"❌ 嗅探错误: {e}")
        else:
            logger.info("🔄 嗅探暂停，等待重启指令...")
            time.sleep(1)
    logger.info("🛑 嗅探线程退出")


def handle_wifi_connection(wifi_string, sniffing_event, exit_event):
    """
    处理WiFi连接流程（包含模式切换、连接和错误恢复）
    
    参数:
        target_data: 收集到的凭证数据列表
        sniffing_event: 线程事件控制嗅探状态
        exit_event: 线程事件控制程序退出
    """

    try:
        # 2. 解析凭证
        ssid, password = parse_wifi_credentials(wifi_string)
        logger.info(f"🔑 尝试连接: SSID='{ssid}'")

        # 3. 执行连接流程
        sniffing_event.clear()  # 暂停嗅探
        if  _establish_wifi_connection(ssid, password):
            broadcast_ack()
            logger.info("✅ 连接成功，程序退出")
            exit_event.set()
        else:
            raise RuntimeError("连接流程失败")

    except subprocess.CalledProcessError as e:
        logger.info(f"❌ 连接失败: {e.stderr.strip()}")
        _recover_to_monitor(sniffing_event, exit_event)

    except ValueError as e:
        logger.info(f"❌ 凭证无效: {e}")
        sniffing_event.set()  # 恢复嗅探

    except Exception as e:
        logger.info(f"‼️ 未知错误: {e}")
        _recover_to_monitor(sniffing_event, exit_event)

    logger.info("=" * 50)

def _establish_wifi_connection(ssid, password, monitor_iface="wlan0mon", managed_iface="wlan0"):
    """执行实际的WiFi连接流程"""
    global INTERFACE

    # 1. 切换到托管模式
    switch_to_managed(monitor_iface, managed_iface)

    # 2. 执行连接命令
    result = subprocess.run(
        ["nmcli", "device", "wifi", "connect", ssid, "password", password, "ifname", managed_iface],
        check=True,
        capture_output=True,
        text=True
    )
    logger.info(f"连接结果: {result.stdout}")
    # 3. 验证连接结果
    if "successfully" not in result.stdout.lower() and "成功" not in result.stdout:
        raise RuntimeError("连接响应异常")
    
    INTERFACE = managed_iface
    return True

def _recover_to_monitor(sniffing_event, exit_event):
    """恢复监控模式并处理异常状态"""
    if not switch_to_monitor():
        logger.info("❌ 无法恢复监控模式!")
        exit_event.set()
    else:
        sniffing_event.set()  # 恢复嗅探


def find_duplicate_string(target_data):
    global has_check_repeat_flag
    """
    检查target_data列表中是否存在重复字符串
    
    参数:
        target_data: 字符串列表
        
    返回:
        str: 返回第一个被发现的重复字符串
        None: 如果没有重复则返回None
    """
    seen = set()
    for item in target_data:
        if item in seen:
            return item  # 返回第一个重复的字符串
        seen.add(item)
    has_check_repeat_flag = True
    return None  # 没有重复


def handle_packet(pkt):
    global is_active, frame_data, data_dict, total_len, total_len_list, total_len_count
    global current_index, target_data, sniffing_event, exit_event, has_check_repeat_flag
    
    if not (pkt.haslayer(Dot11) and pkt.type == 2 and pkt.addr1 == BROADCAST_MAC):
        return
    
    frame_len = len(raw(pkt))
    if frame_len < HEADER_LEN:
        return
    
    ascii_code = frame_len - HEADER_LEN
    timestamp = time.strftime("%H:%M:%S", time.localtime())
    logger.info(f"\n🌐 [{timestamp}] 捕获帧 (长度: {ascii_code} 字节)")
    
    frame_data.append(ascii_code)
    frame_data_len = len(frame_data)
    if ascii_code == ASCII_START:
        if frame_data_len > 1:
            # 开始帧不在首位，帧读取错误，清空数据，以当前开始帧重新开始
            reset()
            frame_data.append(ascii_code)
        is_active = True
        return
    
    if ascii_code == ASCII_END:
        logger.info("  🛑 检测到结束帧")
        logger.info(f"  📋 帧数据: {frame_data}")
        logger.info(f"  📝 解码结果: {data_dict}")
        logger.info(f"  🔢 总长度计数: {total_len_count}")
        logger.info(f"  📊 总长度候选: {total_len_list}")
        logger.info(f"  📏 最终总长度: {total_len}")
        logger.info(f"  🎉 解码结果集合: {target_data}")
        logger.info('-' * 50)
        reset()
        return
    
    if not is_active:
        return
    
    if len(frame_data) > 5:
        logger.info("  ⚠️ 帧数量过多，重置")
        reset()
        return
    
    if len(frame_data) == 2:
        temp_len = ascii_code
        logger.info(f"  📏 当前总长度: {temp_len}")
        if total_len_count < 5:
            total_len_list.append(temp_len)
            total_len_count += 1
        elif total_len_count >= 5 and total_len == 0:
            total_len = max(set(total_len_list), key=total_len_list.count)
            total_len_list.clear()
            logger.info(f"  ✅ 确认总长度: {total_len}")
    
    elif len(frame_data) == 3:
        current_index = ascii_code
    
    elif len(frame_data) == 4:
        char = ascii_char(ascii_code)
        logger.info(f"  🔤 解码字符: {char}")
        if current_index >= 0 and len(char) == 1:
            data_dict[current_index] = char
            logger.info(f"data_dict: {data_dict}")
        else:
            logger.info("  ❌ 解码失败，丢弃帧")
    
    if total_len > 0:
        if len(data_dict) == total_len:
            result = "".join(data_dict[k] for k in sorted(data_dict.keys()))
            logger.info(f"  🎉 解码完成: {result}")
            target_data.append(result)
            has_check_repeat_flag = False
            data_dict.clear()
            reset()
        elif len(data_dict) > total_len:
            result = "".join(data_dict[k] for k in sorted(data_dict.keys())[:total_len])
            logger.info(f"  🎉 存在多余数据，进行截取，解码完成: {result}")
            target_data.append(result)
            has_check_repeat_flag = False
            data_dict.clear()
            reset()

    if len(target_data) > 1 and not has_check_repeat_flag:
        # 如果存在两个一样的字符串，则直接发送这个字符串
        wifi_string = find_duplicate_string(target_data)
        if wifi_string:
            logger.info("快捷算法：有两个相同的字符串直接发送！")
            target_data.clear()
            handle_wifi_connection(wifi_string, sniffing_event, exit_event)

    if len(target_data) > 5:
        # 1. 提取出现频率最高的凭证
        wifi_string = max(set(target_data), key=target_data.count)
        logger.info('*' * 50)
        logger.info(f"  🎉 最佳候选凭证: {wifi_string}")
        logger.info('*' * 50)
        target_data.clear()
        handle_wifi_connection(wifi_string, sniffing_event, exit_event)

    logger.info("=" * 50)


def set_channel(iface, channel):
    """设置无线接口到指定信道"""
    try:
        logger.info(f"🔧 设置接口 {iface} 到信道 {channel}")
        subprocess.run(
            ["iw", iface, "set", "channel", str(channel)],
            check=True,
            capture_output=True,
            text=True
        )
        time.sleep(1)  # 等待信道切换稳定
        return True
    except subprocess.CalledProcessError as e:
        logger.info(f"❌ 设置信道 {channel} 失败: {e.stderr.strip()}")
        return False

def scan_channels(iface="wlan0mon", channels=[1, 6, 11], duration=1):
    """扫描指定信道，统计每个信道的包数量，并将包最多的信道设置为当前信道"""
    packet_counts = {ch: 0 for ch in channels}
    
    def packet_counter(pkt):
        """回调函数，统计数据包数量"""
        if pkt.haslayer(Dot11):
            packet_counts[current_channel] += 1
    
    for channel in channels:
        if not set_channel(iface, channel):
            logger.info(f"❌ 跳过信道 {channel} 的扫描")
            continue
        
        logger.info(f"📡 开始扫描信道 {channel}，持续 {duration} 秒")
        try:
            conf.use_pcap = True  # 确保使用 libpcap
            current_channel = channel
            packet_counts[channel] = 0
            sniff(iface=iface, prn=packet_counter, store=0, monitor=True, timeout=duration)
            logger.info(f"✅ 信道 {channel} 扫描完成，接收到 {packet_counts[channel]} 个数据包")
        except Exception as e:
            logger.info(f"❌ 信道 {channel} 扫描错误: {e}")
    
    # 选择包数量最多的信道
    if packet_counts:
        best_channel = max(packet_counts, key=packet_counts.get)
        max_packets = packet_counts[best_channel]
        logger.info(f"\n🔍 最佳信道: {best_channel} (接收到 {max_packets} 个数据包)")
        
        # 设置为最佳信道
        if set_channel(iface, best_channel):
            logger.info(f"✅ 接口 {iface} 已设置为信道 {best_channel}")
        else:
            logger.info(f"❌ 无法将接口 {iface} 设置为信道 {best_channel}，保持当前信道")
    else:
        logger.info("\n❌ 所有信道扫描失败，无法选择最佳信道")
        best_channel = None
    
    return packet_counts, best_channel

def main():
    check_root()
    logger.info(f"🎯 目标MAC: {BROADCAST_MAC}")
    logger.info("按 Ctrl+C 停止")
    
    global sniffing_thread, sniffing_event, exit_event
    try:
        if is_wifi_connected():
            logger.info("当前已连接无线网络")
            sys.exit(0)
        else:
            logger.info("未连接无线网络")
        
        switch_to_monitor()
        # use_pcap非常重要，启用使用libcap比起原始套接字更兼容和稳定！如果不配置，会导致无法重启嗅探！
        conf.use_pcap = True
        # 扫描信道，确定发信源信道
        scan_channels()
        # 正式嗅探捕获数据
        sniffing_event.set()
        sniffing_thread = threading.Thread(target=start_sniffing, daemon=True)
        sniffing_thread.start()
        while not exit_event.is_set():  # 监控退出信号
            time.sleep(1)
        logger.info("🛑 主线程退出")
    except KeyboardInterrupt:
        logger.info("\n🛑 用户停止监听")
        exit_event.set()



logger = setup_logger()

if __name__ == "__main__":
    main()